# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
### Reducing Prejudice and Support for Religious 
### Nationalism Through Conversations on WhatsApp 

### Author: Rajeshwari Majumdar
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# File:     setup_functions.R
# Purpose:  create custom functions for model estimation and output 
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

# Main text regressions ---- 
fn_reg_main = function(depvar, data, topic = "", print = "short") {
  data_hindus = data %>% filter(religion == "Hindu")
  
  formula = as.formula(paste(depvar, "~ condition_pair + prejudice_binary"))
  
  model.all = lm(formula, data = data_hindus)
  model.non = lm(formula, data = data_hindus %>% filter(condition_topic == "non"))
  model.gen = lm(formula, data = data_hindus %>% filter(condition_topic == "gen"))
  model.igr = lm(formula, data = data_hindus %>% filter(condition_topic == "igr"))
  
  if (topic != "all") {
    models = list(model.all, model.non, model.gen, model.igr)
    names(models) = c("All", "Non-pol.", "Gen pol.", "IGR")
  } else {
    models = list(model.all)
    names(models) = "All"
  }
  
  if (print == "short"){
    m = modelsummary(
      models,
      coef_map = c(
        "condition_pairHM" = "Hindu-Muslim pair",
        "prejudice_binary" = "Baseline prejudice",
        "(Intercept)" = "Constant"
      ),
      stars = c(`***` = 0.01, `**` = 0.05, `*` = 0.10),
      gof_map = "nobs"
    ) 
  } else if (print == "full"){
    m = modelsummary(
      models,
      stars = c(`***` = 0.01, `**` = 0.05, `*` = 0.10),
      gof_map = "nobs"
    ) 
  }

  return(m)
}


# Regression tables with more dependent/control variables ----
fn_reg_more = function(depvars, add.controls = NULL, who = "Hindus", treat = "condition_pair", print = "short", data) {
  if (who == "Hindus") {
    data = data %>% filter(religion == "Hindu")
  } else if (who == "Muslims"){
    data = data %>% filter(religion == "Muslim")
  }
  
  models =  lapply(depvars, function(depvar) {
    formula = if (length(add.controls) == 0) {
      as.formula(paste(depvar, "~", treat))
    } else {
      as.formula(paste(depvar, "~", treat, "+", paste(add.controls, collapse = " + ")))
    }
    model = lm(formula, data = data)
    return(model)
  })
  
  if (print == "short"){
    m = modelsummary(
      models,
      coef_map = c(
        "condition_pairHM" = "Hindu-Muslim pair",
        "prejudice_binary" = "Baseline prejudice",
        "(Intercept)" = "Constant"
      ),
      stars = c(`***` = 0.01, `**` = 0.05, `*` = 0.10),
      gof_map = "nobs"
    ) 
  } else if (print == "full"){
    m = modelsummary(
      models,
      stars = c(`***` = 0.01, `**` = 0.05, `*` = 0.10),
      gof_map = "nobs"
    ) 
  }
  
  return(m)
}


# Six-level treatment variable regressions ----
fn_reg_sixer = function(depvars, depvars_names, data){
  data_hindus = data %>% filter(religion == "Hindu")
  
  add.vars = c( 
    # Individual demographics
    "age","gender_f","college","incomeamnt",
    "castefw","state_zonal","married",
    # Individual preferences and attitudes
    "socialapps_WhatsApp", 
    "religiosity_score","supportBJP",
    "prejudice_binary",
    # Treatment month
    "Month"
  )
  
  models = list()
  for (i in 1:length(depvars)){
    formula.base = as.formula(paste0(depvars[i], "~ condition_pt + prejudice_binary"))
    formula.controls = as.formula(paste0(depvars[i], "~ condition_pt + prejudice_binary +", paste(add.vars, collapse = "+")))
    model.base = lm(formula.base, data = data_hindus)
    model.controls = lm(formula.controls, data = data_hindus)
    models[[i]] = list(model.base, model.controls)
  }
  models = unlist(models, recursive = FALSE)
  names(models) =  rep(depvars_names, each = 2)
  
  m = modelsummary(
    models,
    add_rows = data.frame(
      term = "Controls",
      value = t(rep(c("No","Yes"), length(depvars)))
    ),
    coef_map = c(
      "condition_ptHH_gen" = "Hindu-Hindu Gen. pol.",
      "condition_ptHH_igr" = "Hindu-Hindu IGR",
      "condition_ptHM_non" = "Hindu-Muslim Non-pol.",
      "condition_ptHM_gen" = "Hindu-Muslim Gen. pol.",
      "condition_ptHM_igr" = "Hindu-Muslim IGR",
      "prejudice_binary" = "Baseline prejudice",
      "(Intercept)" = "Constant"
    ),
    stars = c(`***` = 0.01, `**` = 0.05, `*` = 0.10),
    gof_map = "nobs"
  )
  
  return(m)
}


# Effect sizes (Cohen's d) ----
fn_effectsize = function(depvar, topic, data){
  return(
    if (topic == "all"){
      cohen.d(
        data %>% filter(religion == "Hindu" & condition_pair == "HM") %>% 
          pull(depvar) %>% na.omit(), 
        data %>% filter(religion == "Hindu" & condition_pair == "HH") %>% 
          pull(depvar) %>% na.omit() 
      )
    }
    else if (topic != "all"){
      cohen.d(
        data %>% filter(religion == "Hindu" & condition_pt == paste0("HM_",topic)) %>% 
          pull(depvar) %>% na.omit(),
        data %>% filter(religion == "Hindu" & condition_pt == paste0("HH_",topic)) %>% 
          pull(depvar) %>% na.omit()
      )
    }
  )
}


# Extract coefficients for plots ----
fn_extract_coefs = function(depvars, depvar_labels, timepoints="", controls="prejudice_binary", topic="all", who="Hindus", scalevars=FALSE, data){
  if (who == "Hindus"){data = data %>% filter(religion == "Hindu")}
  
  if (topic != "all"){data = data %>% filter(condition_topic == topic)}
  
  input.df = data.frame(depvar = depvars, depvar_label = depvar_labels, timepoint = timepoints, control = controls, topic = topic)
  
  results = bind_rows(lapply(seq_along(input.df$depvar), function(i) {
    if (scalevars == TRUE) {
      scaled_name = paste0(input.df$depvar[i], "_scaled")
      data[[scaled_name]] = scale(data[[input.df$depvar[i]]])
      input.df$depvar[i] = scaled_name
    }
    formula = as.formula(paste0(input.df$depvar[i], "~ condition_pair +", input.df$control[i]))
    mod = lm(formula, data = data)
    results = broom::tidy(mod)
    results = results %>% 
      filter(term == "condition_pairHM") %>% 
      mutate(
        depvar = input.df$depvar[i],
        depvar_label = input.df$depvar_label[i],
        timepoint = input.df$timepoint[i], 
        control = input.df$control[i],
        topic = input.df$topic[i]
      )
    results$depvar = sub("_scaled$", "", results$depvar)
    return(results)
  }))
}
